package icu.zhhll.wrapper;

import icu.zhhll.utils.HttpHelper;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

/**
 * 可重复读取body体
 * <p>
 * 原始Request只支持读取一次body体内容
 * 使用BodyReaderHttpServletRequestWrapper将request进行包装
 * BodyReaderHttpServletRequestWrapper requestWrapper = new BodyReaderHttpServletRequestWrapper(request);
 * 之后获取body体内容使用requestWrapper
 * ServletInputStream inputStream = requestWrapper.getInputStream();
 * 这样可使得body体重复读取
 *
 * @author zh
 * @date 2023/8/14 17:13
 */
public class BodyReaderHttpServletRequestWrapper extends HttpServletRequestWrapper {
    private final ThreadLocal<Map<String, String>> headLocal = ThreadLocal.withInitial(HashMap::new);

    /**
     * 删除的请求头只会取用本地线程存储中的
     */
    private final ThreadLocal<List<String>> removeLocal = ThreadLocal.withInitial(ArrayList::new);

    private byte[] body;

//	private InputStream input;

    private final HttpServletRequest orgRequest;

    public BodyReaderHttpServletRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
        this.orgRequest = request;
        body = HttpHelper.getBody(request);
    }

    public HttpServletRequest getOrgRequest() {
        return this.orgRequest;
    }

    @Override
    public BufferedReader getReader() {
        return new BufferedReader(new InputStreamReader(getInputStream()));
    }

    public void setInput(byte[] body) {
        this.body = body;
    }

    /**
     * 添加请求头
     *
     * @param name
     * @param value
     */
    public void addHeader(String name, String value) {
        if (headLocal.get() != null) {
            headLocal.get().put(name, value);
        }
    }

    /**
     * 删除请求头
     *
     * @param name
     */
    public void removeHeader(String name) {
        if (headLocal.get() != null) {
            headLocal.get().remove(name);
        }
        removeLocal.get().add(name);
    }

    /**
     * 获取请求头
     *
     * @param name
     * @return
     */
    @Override
    public String getHeader(String name) {
        // 如果本地线程中包含该头信息  取本地线程存储中的
        if (headLocal.get() != null && headLocal.get().containsKey(name)) {
            return headLocal.get().get(name);
        }
        if (removeLocal.get().contains(name)) {
            return null;
        }
        // 如果没有取用原始请求的
        return this.orgRequest.getHeader(name);
    }

    /**
     * 获取请求头
     *
     * @param name
     * @return
     */
    @Override
    public Enumeration<String> getHeaders(String name) {
        if (headLocal.get() != null && headLocal.get().containsKey(name)) {
            Vector<String> vector = new Vector<>();
            vector.add(headLocal.get().get(name));
            if (!removeLocal.get().contains(name) && this.orgRequest.getHeaders(name) != null) {
                Enumeration<String> headers = this.orgRequest.getHeaders(name);
                while (headers.hasMoreElements()) {
                    vector.addElement(headers.nextElement());
                }
            }
            return vector.elements();
        }
        if (removeLocal.get().contains(name)) {
            return null;
        }
        return this.orgRequest.getHeaders(name);
    }

    /**
     * 获取请求头名字
     *
     * @return
     */
    @Override
    public Enumeration<String> getHeaderNames() {
        if (headLocal.get() != null && headLocal.get().size() > 0) {
            Vector<String> vector = new Vector<>(headLocal.get().keySet());
            Enumeration<String> headers = this.orgRequest.getHeaderNames();
            if (headers != null) {
                while (headers.hasMoreElements()) {
                    String element = headers.nextElement();
                    if (removeLocal.get().contains(element)) {
                        continue;
                    }
                    vector.addElement(element);
                }
            }
            return vector.elements();
        } else if (removeLocal.get().isEmpty()) {
            return this.orgRequest.getHeaderNames();
        } else {
            Vector<String> vector = new Vector<>();
            Enumeration<String> headers = this.orgRequest.getHeaderNames();
            if (headers != null) {
                while (headers.hasMoreElements()) {
                    String element = headers.nextElement();
                    if (removeLocal.get().contains(element)) {
                        continue;
                    }
                    vector.addElement(element);
                }
            }
            return vector.elements();
        }
    }


    @Override
    public ServletInputStream getInputStream() {


        final ByteArrayInputStream bais = new ByteArrayInputStream(body);

        return new ServletInputStream() {

            @Override
            public int read() {
                return bais.read();
            }

            @Override
            public boolean isFinished() {
                return false;
            }

            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setReadListener(ReadListener readListener) {

            }
        };
    }
}
